import hashlib
import hmac
import os
import json
import time
import base64
from typing import Tuple, Union

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.exceptions import InvalidTag  # 用于捕获GCM解密认证失败


# --- 1. 推荐的密钥设计方案 ---

# 1.1 密钥派生函数 (与你提供的示例一致)
def derive_aes_key(api_key: str, api_secret: str, salt: str = "fixed_salt", iterations: int = 100000,
                   key_length: int = 32) -> bytes:
    """
    基于api_key, api_secret和salt派生AES密钥。这是耗时操作，客户端自己派生，服务器最好保存起来，不用每次都计算，属于计算密集型
    """
    # 注意：你提供的示例中 key_material = api_key + api_secret + salt
    # 然后 salt.encode() 又作为 salt 参数传递给 pbkdf2_hmac。
    # 这意味着 salt 字符串实际上被使用了两次。这里我们遵循你的示例代码。
    key_material = (api_key + api_secret + salt).encode('utf-8')
    derived_key = hashlib.pbkdf2_hmac(
        'sha256',
        key_material,  # password
        salt.encode('utf-8'),  # salt
        iterations,
        dklen=key_length  # AES-256 需要 32 字节的密钥
    )
    return derived_key


# --- 服务端模拟存储 ---
# 在实际应用中，这会是一个安全的数据库
SERVER_USER_STORE = {}
FIXED_SALT_FOR_AES_DERIVATION = "your_application_wide_fixed_salt_for_aes_derivation"


def register_user_on_server(api_key: str, api_secret: str):
    """模拟在服务器端注册用户并存储其凭证和派生的AES密钥"""
    if api_key in SERVER_USER_STORE:
        print(f"用户 {api_key} 已存在。")
        return

    derived_key_bytes = derive_aes_key(api_key, api_secret, FIXED_SALT_FOR_AES_DERIVATION)
    derived_key_hex = derived_key_bytes.hex()  # <--- 转换为hex

    SERVER_USER_STORE[api_key] = {
        "api_secret": api_secret,
        "derived_aes_key_hex": derived_key_hex  # <--- 存储hex字符串
    }
    print(f"用户 {api_key} 注册成功，AES密钥已派生并以hex形式存储。")


# --- 2. 完整的加密方案建议 ---

# AES-GCM 相关参数
AES_KEY_SIZE_BYTES = 32  # 256位
GCM_IV_LENGTH_BYTES = 12  # 96位 (12字节) 是GCM推荐的IV长度
GCM_TAG_LENGTH_BYTES = 16  # 128位 (16字节) 认证标签


# 2.1 AES-GCM 加密函数
def encrypt_aes_gcm(plaintext: str, aes_key: bytes) -> Tuple[bytes, bytes, bytes]:
    """
    使用AES-GCM加密数据。
    返回: (ciphertext, iv, tag)
    """
    if len(aes_key) != AES_KEY_SIZE_BYTES:
        raise ValueError(f"AES密钥长度必须为 {AES_KEY_SIZE_BYTES} 字节")

    iv = os.urandom(GCM_IV_LENGTH_BYTES)
    cipher = Cipher(algorithms.AES(aes_key), modes.GCM(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    # GCM模式下，可以有“附加认证数据”(AAD)，这里我们不使用，但可以按需添加
    # encryptor.authenticate_additional_data(b"some_associated_data")
    ciphertext = encryptor.update(plaintext.encode('utf-8')) + encryptor.finalize()
    return ciphertext, iv, encryptor.tag


# 2.2 AES-GCM 解密函数
def decrypt_aes_gcm(ciphertext: bytes, aes_key: bytes, iv: bytes, tag: bytes) -> Union[str, None]:
    """
    使用AES-GCM解密数据。
    如果认证失败 (tag不匹配或数据被篡改)，会抛出 InvalidTag 异常。
    返回: 解密后的明文字符串，或在解密失败时返回 None
    """
    if len(aes_key) != AES_KEY_SIZE_BYTES:
        raise ValueError(f"AES密钥长度必须为 {AES_KEY_SIZE_BYTES} 字节")
    if len(iv) != GCM_IV_LENGTH_BYTES:
        raise ValueError(f"IV长度必须为 {GCM_IV_LENGTH_BYTES} 字节")

    cipher = Cipher(algorithms.AES(aes_key), modes.GCM(iv, tag), backend=default_backend())
    decryptor = cipher.decryptor()
    # 如果加密时使用了AAD，解密时也必须使用相同的AAD
    # decryptor.authenticate_additional_data(b"some_associated_data")
    try:
        plaintext_bytes = decryptor.update(ciphertext) + decryptor.finalize()
        return plaintext_bytes.decode('utf-8')
    except InvalidTag:
        print("AES-GCM 解密失败：认证标签无效！数据可能被篡改或密钥/IV错误。")
        return None
    except Exception as e:
        print(f"AES-GCM 解密时发生其他错误: {e}")
        return None


# 2.3 HMAC签名生成函数
def generate_request_signature(api_secret: str, data_to_sign_str: str) -> str:
    """
    使用HMAC-SHA256为请求数据生成签名。
    data_to_sign_str: 需要签名的请求数据的规范化字符串表示。
    """
    message = data_to_sign_str.encode('utf-8')
    secret = api_secret.encode('utf-8')
    signature = hmac.new(secret, message, hashlib.sha256).hexdigest()
    return signature


# 2.4 HMAC签名验证函数
def verify_request_signature(api_secret: str, received_signature: str, data_to_sign_str: str) -> bool:
    """
    验证接收到的签名是否有效。
    """
    expected_signature = generate_request_signature(api_secret, data_to_sign_str)
    return hmac.compare_digest(expected_signature, received_signature)


# --- 模拟客户端和服务端交互 ---

def client_side_process(api_key: str, api_secret: str, sensitive_payload: dict):
    """模拟客户端处理流程"""
    print("\n--- 客户端开始处理 ---")

    # 1. 客户端派生AES密钥 (客户端不需要存储，每次请求时派生)
    #    注意：客户端使用的 salt 必须与服务端派生时使用的 salt 一致
    client_derived_aes_key = derive_aes_key(api_key, api_secret, FIXED_SALT_FOR_AES_DERIVATION)
    print(f"客户端为 {api_key} 派生的AES密钥 (前8字节): {client_derived_aes_key[:8].hex()}...")

    # 2. 加密敏感数据
    plaintext_data_to_encrypt = json.dumps(sensitive_payload)
    ciphertext, iv, tag = encrypt_aes_gcm(plaintext_data_to_encrypt, client_derived_aes_key)

    # 为了在JSON中传输，通常将二进制数据进行Base64编码
    encrypted_data_b64 = base64.b64encode(ciphertext).decode('utf-8')
    iv_b64 = base64.b64encode(iv).decode('utf-8')
    tag_b64 = base64.b64encode(tag).decode('utf-8')  # GCM tag也需要传输

    print(f"原始敏感数据: {plaintext_data_to_encrypt}")
    print(f"加密后的数据 (Base64): {encrypted_data_b64[:30]}...")
    print(f"IV (Base64): {iv_b64}")
    print(f"Tag (Base64): {tag_b64}")

    # 3. 准备请求体
    timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    request_body_before_signing = {
        "api_key": api_key,
        "timestamp": timestamp,
        # "encrypted_data" 和 "iv" 也是签名的一部分
        # 注意：签名的数据应该是一个规范化的字符串，以确保客户端和服务端生成一致的签名源。
        # 这里我们简单地将要包含在签名内的字段按key排序后拼接。
        # 实际应用中，更推荐对一个包含所有待签名部分的JSON对象的规范化字符串进行签名。
    }

    # 4. 生成请求签名
    # 定义哪些字段参与签名。这里我们选择 api_key, timestamp, encrypted_data, iv。
    # 重要的是，客户端和服务端必须对完全相同的内容进行签名。
    # 顺序也很重要，所以通常会对字段名排序或使用固定的顺序。
    # 为了简单，这里我们构造一个字符串。
    # 注意：这里我们不直接使用整个 request_body_before_signing 的 JSON 字符串，
    # 因为 "signature" 字段本身不应包含在被签名的内容中。
    # 我们将加密数据和 IV 也作为签名内容的一部分。
    data_parts_for_signing = [
        ("api_key", api_key),
        ("timestamp", timestamp),
        ("encrypted_data", encrypted_data_b64),  # 使用 Base64 编码后的数据
        ("iv", iv_b64)  # 使用 Base64 编码后的 IV
    ]
    # 按key排序以确保一致性
    # data_parts_for_signing.sort(key=lambda x: x[0]) # 如果字段较多，排序是个好习惯
    # 构造规范化字符串
    # canonical_string_to_sign = "&".join([f"{k}={v}" for k, v in data_parts_for_signing])
    # 或者更简单地，按固定顺序拼接 (如你示例中的 request_data)
    # 假设 request_data 是指明文部分的某些字段和加密后的数据。
    # 你的示例是 HMAC-SHA256(api_secret, request_data)
    # 这里的 request_data 可以是 timestamp + encrypted_data_b64 + iv_b64
    # 让我们更明确一点：签名内容应该是请求中除了签名本身之外，能被服务端验证的部分。
    # 通常是：HTTP方法 + URI + Query参数 + Timestamp + 部分Body内容（如api_key, encrypted_data）。
    # 为了此示例，我们简化为：
    string_to_sign = f"{api_key}{timestamp}{encrypted_data_b64}{iv_b64}"
    signature = generate_request_signature(api_secret, string_to_sign)
    print(f"用于签名的字符串: {string_to_sign[:50]}...")
    print(f"生成的签名: {signature}")

    # 5. 构建最终请求 (模拟JSON请求)
    final_request = {
        "api_key": api_key,
        "timestamp": timestamp,
        "signature": signature,
        "encrypted_data": encrypted_data_b64,
        "iv": iv_b64,
        "tag": tag_b64  # GCM tag也需要传递给服务器
    }
    print(f"最终发送的请求 (JSON): {json.dumps(final_request, indent=2)}")
    return final_request


def server_side_process(request: dict):
    """模拟服务端处理流程"""
    print("\n--- 服务端开始处理 ---")
    print(f"接收到的请求: {json.dumps(request, indent=2)}")

    api_key = request.get("api_key")
    received_signature = request.get("signature")
    timestamp = request.get("timestamp")
    encrypted_data_b64 = request.get("encrypted_data")
    iv_b64 = request.get("iv")
    tag_b64 = request.get("tag")  # 接收GCM tag

    # 1. 验证api_key并获取用户信息
    user_info = SERVER_USER_STORE.get(api_key)
    if not user_info:
        print(f"错误：无效的 api_key: {api_key}")
        return

    api_secret = user_info["api_secret"]
    # 从存储中获取hex字符串并转换回bytes
    derived_aes_key_hex_from_store = user_info["derived_aes_key_hex"]
    server_derived_aes_key = bytes.fromhex(derived_aes_key_hex_from_store)  # <--- 从hex转换回bytes

    print(f"服务端为 {api_key} 使用的AES密钥 (前8字节的hex): {server_derived_aes_key[:8].hex()}...")

    # 2. 验证签名
    # 服务端必须以与客户端完全相同的方式构造待签名的字符串
    string_to_verify = f"{api_key}{timestamp}{encrypted_data_b64}{iv_b64}"
    print(f"服务端重构的用于验证签名的字符串: {string_to_verify[:50]}...")
    if not verify_request_signature(api_secret, received_signature, string_to_verify):
        print("错误：签名验证失败！")
        return
    print("签名验证成功！")

    # （可选）验证时间戳，防止重放攻击
    current_time = time.mktime(time.gmtime())
    request_time = time.mktime(time.strptime(timestamp, "%Y-%m-%dT%H:%M:%SZ"))
    if abs(current_time - request_time) > 300:  # 例如，5分钟窗口
        print("警告：时间戳可能已过期（或时钟不同步）。")
        # 根据策略决定是否拒绝

    # 3. 解密数据
    try:
        ciphertext = base64.b64decode(encrypted_data_b64)
        iv = base64.b64decode(iv_b64)
        tag = base64.b64decode(tag_b64)  # Base64解码GCM tag

        decrypted_payload_str = decrypt_aes_gcm(ciphertext, server_derived_aes_key, iv, tag)

        if decrypted_payload_str:
            decrypted_payload = json.loads(decrypted_payload_str)
            print(f"解密后的敏感数据: {decrypted_payload}")
            # 在这里处理解密后的业务逻辑...
            print("服务端处理完成。")
        else:
            print("服务端错误：数据解密失败（认证未通过或数据损坏）。")

    except Exception as e:
        print(f"服务端解密或处理时发生错误: {e}")


# --- 主程序执行 ---
if __name__ == "__main__":
    # 0. 初始化：在服务端注册两个用户
    register_user_on_server("user123_apikey", "user123_apisecret_very_secure")
    register_user_on_server("user456_apikey", "user456_apisecret_super_safe")
    print(f"\n当前服务端用户存储: {list(SERVER_USER_STORE.keys())}")

    # 模拟用户1的请求
    client1_api_key = "user123_apikey"
    client1_api_secret = "user123_apisecret_very_secure"  # 客户端知道自己的secret
    client1_payload = {"action": "transfer_funds", "amount": 100, "to_account": "ACC789"}

    simulated_request_from_client1 = client_side_process(client1_api_key, client1_api_secret, client1_payload)
    server_side_process(simulated_request_from_client1)

    print("\n" + "=" * 50 + "\n")

    # 模拟用户2的请求
    client2_api_key = "user456_apikey"
    client2_api_secret = "user456_apisecret_super_safe"
    client2_payload = {"action": "get_balance", "account_id": "ACCXYZ"}

    simulated_request_from_client2 = client_side_process(client2_api_key, client2_api_secret, client2_payload)
    server_side_process(simulated_request_from_client2)

    print("\n" + "=" * 50 + "\n")
    # 模拟一个签名篡改的场景
    print("--- 模拟签名篡改场景 ---")
    client3_api_key = "user123_apikey"
    client3_api_secret = "user123_apisecret_very_secure"
    client3_payload = {"action": "delete_user", "user_id_to_delete": "unwanted_user"}

    request_before_tampering = client_side_process(client3_api_key, client3_api_secret, client3_payload)

    # 攻击者篡改签名
    tampered_request = request_before_tampering.copy()
    original_signature = tampered_request["signature"]
    tampered_request["signature"] = original_signature[:-5] + "abcde"  # 随意修改签名
    print(f"原始签名: {original_signature}, 篡改后签名: {tampered_request['signature']}")

    server_side_process(tampered_request)

    print("\n" + "=" * 50 + "\n")
    # 模拟一个加密数据篡改的场景 (GCM应该能检测到)
    print("--- 模拟加密数据篡改场景 (GCM的tag应该使其解密失败) ---")
    request_before_data_tampering = client_side_process(client3_api_key, client3_api_secret, {"original_data": "safe"})

    tampered_data_request = request_before_data_tampering.copy()
    original_encrypted_data_b64 = tampered_data_request["encrypted_data"]
    # 尝试修改加密数据 (这很可能会导致Base64解码失败或解密时tag不匹配)
    # 简单地在Base64字符串末尾加一个字符
    if original_encrypted_data_b64.endswith("="):  # 处理padding
        tampered_data_request["encrypted_data"] = original_encrypted_data_b64[:-1] + "A"
    else:
        tampered_data_request["encrypted_data"] = original_encrypted_data_b64 + "A"
    print(
        f"原始加密数据 (Base64): {original_encrypted_data_b64[:20]}..., 篡改后: {tampered_data_request['encrypted_data'][:20]}...")

    server_side_process(tampered_data_request)
